#include "GraphData.h"
#include <cmath>
#include <unordered_map>
#include <cstdlib>
#include <ctime>
#include <limits>
#include <algorithm>
#include <iostream>
#include <fstream>
#include <vector>
#include <string>
#include <algorithm>
#include <random>

namespace {

// 欧氏距离
double euclideanDistance(const std::vector<double>& a, const std::vector<double>& b) {
    double sum = 0.0;
    for (size_t i = 0; i < a.size(); ++i) {
        double d = a[i] - b[i];
        sum += d * d;
    }
    return std::sqrt(sum);
}

// 初始化质心（随机打乱选择前 k 个）
std::vector<std::vector<double>> initializeCentroids(const std::vector<std::vector<double>>& data, int k) {
    std::vector<size_t> indices(data.size());
    std::iota(indices.begin(), indices.end(), 0);  // 生成 0, 1, ..., data.size() - 1

    // 使用现代随机数生成器打乱索引
    std::random_device rd;
    std::mt19937 gen(rd());
    std::shuffle(indices.begin(), indices.end(), gen);

    std::vector<std::vector<double>> centroids;
    for (int i = 0; i < k; ++i) {
        centroids.push_back(data[indices[i]]);
    }

    return centroids;
}

// 核心 KMeans 实现
std::vector<int> kMeans(const std::vector<std::vector<double>>& data, int k, int maxIterations) {
    size_t dim = data[0].size();
    auto centroids = initializeCentroids(data, k);
    std::vector<int> assignments(data.size(), 0);

    for (int iter = 0; iter < maxIterations; ++iter) {
        bool changed = false;

        // 分配阶段
        for (size_t i = 0; i < data.size(); ++i) {
            double minDist = std::numeric_limits<double>::max();
            int minIndex = 0;
            for (int j = 0; j < k; ++j) {
                double dist = euclideanDistance(data[i], centroids[j]);
                if (dist < minDist) {
                    minDist = dist;
                    minIndex = j;
                }
            }
            if (assignments[i] != minIndex) {
                assignments[i] = minIndex;
                changed = true;
            }
        }

        if (!changed) break;

        // 更新阶段
        std::vector<std::vector<double>> sums(k, std::vector<double>(dim, 0.0));
        std::vector<int> counts(k, 0);

        for (size_t i = 0; i < data.size(); ++i) {
            int cluster = assignments[i];
            for (size_t d = 0; d < dim; ++d) {
                sums[cluster][d] += data[i][d];
            }
            counts[cluster]++;
        }

        for (int j = 0; j < k; ++j) {
            if (counts[j] == 0) continue;
            for (size_t d = 0; d < dim; ++d) {
                centroids[j][d] = sums[j][d] / counts[j];
            }
        }
    }

    return assignments;
}

} // namespace

// clusterEdges 实现
ClusteredEdges clusterEdges(const std::vector<Edge>& edges,
                            const std::vector<Node>& nodes,
                            int k,
                            int maxIterations) {
    std::unordered_map<int, Node> nodeMap;
    for (const auto& node : nodes) {
        nodeMap[node.id] = node;
    }

    std::vector<std::vector<double>> features;
    for (const auto& edge : edges) {
        const Node& p1 = nodeMap.at(edge.from);
        const Node& p2 = nodeMap.at(edge.to);

        double dx = p2.lon - p1.lon;
        double dy = p2.lat - p1.lat;
        double len = std::sqrt(dx * dx + dy * dy);
        if (len < 1e-6) len = 1e-6;

        features.push_back({
            dx / len, dy / len,
            (p1.lon + p2.lon) / 2.0,
            (p1.lat + p2.lat) / 2.0
        });
    }

    auto assignments = kMeans(features, k, maxIterations);
    ClusteredEdges result(k);
    for (size_t i = 0; i < edges.size(); ++i) {
        result[assignments[i]].push_back(edges[i]);
    }

    return result;
}


/**
 * 写入聚类后的边结果到 CSV 文件。
 * @param clusteredEdges 每个子向量代表一个聚类组，内部是 (from, to) 的边
 * @param outputPath 输出文件路径
 */
void write_clustered_edges_to_csv(const std::vector<std::vector<Edge>>& clusters, const std::string& filename) {
    std::ofstream outFile(filename);
    if (!outFile.is_open()) {
        throw std::runtime_error("Failed to open output file: " + filename);
    }

    outFile << "cluster,from,to\n";

    for (size_t clusterIdx = 0; clusterIdx < clusters.size(); ++clusterIdx) {
        for (const auto& edge : clusters[clusterIdx]) {
            outFile << clusterIdx << "," << edge.from << "," << edge.to << "\n";
        }
    }

    outFile.close();
}